#include "lua.h"
#include "lauxlib.h"
#include "cm_mqtt.h"
#include "lua_test_entry.h"
#include <string.h>
#include <stdlib.h>

#define LUA_MQTT_CLIENT_METATABLE "mqtt_client_mt"
#define MAX_MQTT_CLIENTS 5

typedef struct {
    cm_mqtt_client_t *client;
    int self_ref;
    int on_connect_ref;
    int on_disconnect_ref;
    int on_message_ref;
} lua_mqtt_client_t;

// Global map to find our Lua client from a C client pointer
static lua_mqtt_client_t* g_client_map[MAX_MQTT_CLIENTS] = {NULL};

static void add_to_client_map(lua_mqtt_client_t* client) {
    for (int i = 0; i < MAX_MQTT_CLIENTS; i++) {
        if (g_client_map[i] == NULL) {
            g_client_map[i] = client;
            return;
        }
    }
}

static void remove_from_client_map(lua_mqtt_client_t* client) {
    for (int i = 0; i < MAX_MQTT_CLIENTS; i++) {
        if (g_client_map[i] == client) {
            g_client_map[i] = NULL;
            return;
        }
    }
}

static lua_mqtt_client_t* find_lua_client(cm_mqtt_client_t* c_client) {
    for (int i = 0; i < MAX_MQTT_CLIENTS; i++) {
        if (g_client_map[i] != NULL && g_client_map[i]->client == c_client) {
            return g_client_map[i];
        }
    }
    return NULL;
}

// --- C Callbacks ---

typedef struct {
    char *topic;
    char *payload;
    int payload_len;
    int self_ref;
} mqtt_message_event_data_t;

static int mqtt_connack_cb(cm_mqtt_client_t* client, int session, cm_mqtt_conn_state_e conn_res) {
    lua_mqtt_client_t *lua_client = find_lua_client(client);
    if (!lua_client) return -1;

    lua_event_t event;
    event.id = (conn_res == CM_MQTT_CONN_STATE_SUCCESS) ? LUA_EVENT_MQTT_CONNACK : LUA_EVENT_MQTT_DISCONNECTED;
    event.data.ival = lua_client->self_ref;
    lua_event_post(&event);
    return 0; // Return int as required by the SDK
}

static int mqtt_publish_cb(cm_mqtt_client_t* client, unsigned short msgid, char* topic, int total_len, int payload_len, char* payload) {
    lua_mqtt_client_t *lua_client = find_lua_client(client);
    if (!lua_client) return 0;

    mqtt_message_event_data_t *msg_data = (mqtt_message_event_data_t*)malloc(sizeof(mqtt_message_event_data_t));
    if (!msg_data) return 0;

    msg_data->topic = strdup(topic);
    msg_data->payload = (char*)malloc(payload_len);
    if (msg_data->payload) {
        memcpy(msg_data->payload, payload, payload_len);
    }
    msg_data->payload_len = payload_len;
    msg_data->self_ref = lua_client->self_ref;

    lua_event_t event;
    event.id = LUA_EVENT_MQTT_MESSAGE;
    event.data.pval = msg_data;
    lua_event_post(&event);

    return 0;
}


// --- Lua Methods ---

static int l_client_on(lua_State *L) {
    lua_mqtt_client_t *lua_client = (lua_mqtt_client_t *)luaL_checkudata(L, 1, LUA_MQTT_CLIENT_METATABLE);
    const char *event_name = luaL_checkstring(L, 2);
    luaL_checktype(L, 3, LUA_TFUNCTION);

    int *target_ref = NULL;
    if (strcmp(event_name, "connect") == 0) target_ref = &lua_client->on_connect_ref;
    else if (strcmp(event_name, "disconnect") == 0) target_ref = &lua_client->on_disconnect_ref;
    else if (strcmp(event_name, "message") == 0) target_ref = &lua_client->on_message_ref;
    else return luaL_argerror(L, 2, "invalid event name");

    if (*target_ref != LUA_NOREF) {
        luaL_unref(L, LUA_REGISTRYINDEX, *target_ref);
    }
    lua_pushvalue(L, 3);
    *target_ref = luaL_ref(L, LUA_REGISTRYINDEX);

    return 0;
}

static int l_client_connect(lua_State *L) {
    lua_mqtt_client_t *lua_client = (lua_mqtt_client_t *)luaL_checkudata(L, 1, LUA_MQTT_CLIENT_METATABLE);
    cm_mqtt_connect_options_t opts = {0};
    opts.hostname = luaL_checkstring(L, 2);
    opts.hostport = luaL_checkinteger(L, 3);
    opts.clientid = luaL_checkstring(L, 4);
    opts.keepalive = luaL_optinteger(L, 5, 120);
    opts.username = luaL_optstring(L, 6, NULL); // Optional username
    opts.password = luaL_optstring(L, 7, NULL); // Optional password
    opts.clean_session = 1;

    int ret = cm_mqtt_client_connect(lua_client->client, &opts);
    lua_pushboolean(L, ret == 0);
    return 1;
}

static int l_client_publish(lua_State *L) {
    lua_mqtt_client_t *lua_client = (lua_mqtt_client_t *)luaL_checkudata(L, 1, LUA_MQTT_CLIENT_METATABLE);
    const char *topic = luaL_checkstring(L, 2);
    size_t payload_len;
    const char *payload = luaL_checklstring(L, 3, &payload_len);
    int qos = luaL_optinteger(L, 4, 0);

    char flags = 0;
    if (qos == 1) flags |= CM_MQTT_QOS_1;
    else if (qos == 2) flags |= CM_MQTT_QOS_2;

    int ret = cm_mqtt_client_publish(lua_client->client, topic, payload, payload_len, flags);
    lua_pushboolean(L, ret >= 0);
    return 1;
}

static int l_client_subscribe(lua_State *L) {
    lua_mqtt_client_t *lua_client = (lua_mqtt_client_t *)luaL_checkudata(L, 1, LUA_MQTT_CLIENT_METATABLE);
    const char *topic = luaL_checkstring(L, 2);
    int qos = luaL_optinteger(L, 3, 0);

    const char* topics[] = {topic};
    char qoses[] = {(char)qos};

    int ret = cm_mqtt_client_subscribe(lua_client->client, topics, qoses, 1);
    lua_pushboolean(L, ret == 0);
    return 1;
}

static int l_client_gc(lua_State *L) {
    lua_mqtt_client_t *lua_client = (lua_mqtt_client_t *)luaL_checkudata(L, 1, LUA_MQTT_CLIENT_METATABLE);

    if (lua_client->client) {
        remove_from_client_map(lua_client);
        cm_mqtt_client_destroy(lua_client->client);
        lua_client->client = NULL;
    }
    if (lua_client->self_ref != LUA_NOREF) luaL_unref(L, LUA_REGISTRYINDEX, lua_client->self_ref);
    if (lua_client->on_connect_ref != LUA_NOREF) luaL_unref(L, LUA_REGISTRYINDEX, lua_client->on_connect_ref);
    if (lua_client->on_disconnect_ref != LUA_NOREF) luaL_unref(L, LUA_REGISTRYINDEX, lua_client->on_disconnect_ref);
    if (lua_client->on_message_ref != LUA_NOREF) luaL_unref(L, LUA_REGISTRYINDEX, lua_client->on_message_ref);

    return 0;
}

static const luaL_Reg client_methods[] = {
    {"on", l_client_on},
    {"connect", l_client_connect},
    {"publish", l_client_publish},
    {"subscribe", l_client_subscribe},
    {"__gc", l_client_gc},
    {NULL, NULL}
};

// Forward declaration
static int l_mqtt_create(lua_State *L);

static const luaL_Reg mqtt_lib[] = {
    {"create", l_mqtt_create},
    {NULL, NULL}
};

int luaopen_mqtt(lua_State *L) {
    luaL_newmetatable(L, LUA_MQTT_CLIENT_METATABLE);
    lua_pushvalue(L, -1);
    lua_setfield(L, -2, "__index");
    luaL_setfuncs(L, client_methods, 0);

    luaL_newlib(L, mqtt_lib);
    return 1;
}

// Definition of l_mqtt_create
static int l_mqtt_create(lua_State *L) {
    lua_mqtt_client_t *lua_client = (lua_mqtt_client_t *)lua_newuserdata(L, sizeof(lua_mqtt_client_t));
    memset(lua_client, 0, sizeof(lua_mqtt_client_t));
    luaL_setmetatable(L, LUA_MQTT_CLIENT_METATABLE);

    lua_client->client = cm_mqtt_client_create();
    if (!lua_client->client) {
        return luaL_error(L, "failed to create cm_mqtt_client");
    }

    lua_client->on_connect_ref = LUA_NOREF;
    lua_client->on_disconnect_ref = LUA_NOREF;
    lua_client->on_message_ref = LUA_NOREF;

    lua_pushvalue(L, -1);
    lua_client->self_ref = luaL_ref(L, LUA_REGISTRYINDEX);

    add_to_client_map(lua_client);

    cm_mqtt_client_cb_t cbs;
    memset(&cbs, 0, sizeof(cm_mqtt_client_cb_t));
    cbs.connack_cb = mqtt_connack_cb;
    cbs.publish_cb = mqtt_publish_cb;
    cm_mqtt_client_set_opt(lua_client->client, CM_MQTT_OPT_EVENT, &cbs);

    return 1;
}
